import uvicorn
from api import roop_fast_api_app, slide_api
from fastapi import Depends
import redis
from config import redis_config
from fastapi import FastAPI, HTTPException
import pickle
from celerytask import app as celery_app
from celery.result import AsyncResult

def get_redis_connection():
    connection = redis.Redis(**redis_config)
    try:
        yield connection
    finally:
        connection.close()


app = FastAPI()
app.include_router(roop_fast_api_app, prefix="/V2", tags=["视频换脸接口"])
app.include_router(slide_api, prefix="/slide", tags=["图片转视频接口"])  # 暂时先关闭这个接口


@app.get("/task/{task_id}")
async def get_status(task_id: str, r: redis.Redis = Depends(get_redis_connection)):
    task = pickle.loads(r.get(task_id)) if r.get(task_id) else None
    if not task:
        raise HTTPException(status_code=404, detail="Task not found")
    return {
        "task_id": task.id,
        "status": task.status,
        "start_time": task.start_time,
        "result": task.result,
        "duration": task.duration
    }


@app.get("/task_status/{task_id}")
async def get_task_status(task_id: str):
    task = AsyncResult(task_id, app=celery_app)
    if task.status == 'SUCCESS':
        return {"status": task.status, "tast_result": task.result}
    return {"status": task.status}


def find_innermost_task_id(result):
    if isinstance(result, list):
        for item in reversed(result):
            if inner_id := find_innermost_task_id(item):
                return inner_id
    elif result:
        return result
    return None


if __name__ == '__main__':
    # 线程有可能会被卡主?
    uvicorn.run("run:app", host='0.0.0.0', port=6006, reload=True, workers=10)
